from keras.models import Sequential
from keras.layers import Conv3D, BatchNormalization, Activation, MaxPooling3D, SpatialDropout3D, Flatten, Dense, Dropout
from keras.regularizers import l2
from resnet3D import ResnetBuilder
# from densenet import densenet_l2
import functools
from vgg13_shortcuts import vgg13_shortcuts, vgg13_shortcuts_v2


def vgg13(input_shape, base_nb_filters, weight_decay):
    model = Sequential()
    model.add(Conv3D(base_nb_filters, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay),
                     input_shape=input_shape))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(SpatialDropout3D(0.25))
    model.add(Conv3D(base_nb_filters*2, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*2, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(SpatialDropout3D(0.25))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(SpatialDropout3D(0.25))
    model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(Flatten())
    model.add(Dropout(0.5))
    model.add(Dense(base_nb_filters*16, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(base_nb_filters*16, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(2, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(Activation('softmax'))

    return model


def vgg10(input_shape, base_nb_filters, weight_decay):
    model = Sequential()
    model.add(Conv3D(base_nb_filters, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay),
                     input_shape=input_shape))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(SpatialDropout3D(0.25))
    model.add(Conv3D(base_nb_filters*2, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*2, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    model.add(SpatialDropout3D(0.25))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv3D(base_nb_filters*4, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))

    print(model.output_shape)

    # model.add(SpatialDropout3D(0.25))
    # model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    # model.add(BatchNormalization())
    # model.add(Activation('relu'))
    # model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    # model.add(BatchNormalization())
    # model.add(Activation('relu'))
    # model.add(Conv3D(base_nb_filters*8, 3, 3, 3, init='he_normal', border_mode='same', W_regularizer=l2(weight_decay)))
    # model.add(BatchNormalization())
    # model.add(Activation('relu'))
    # model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2)))
    #
    # print(model.output_shape)

    model.add(Flatten())
    model.add(Dropout(0.5))
    model.add(Dense(base_nb_filters*16, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(base_nb_filters*16, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(2, init='he_normal', W_regularizer=l2(weight_decay)))
    model.add(Activation('softmax'))

    return model


resnet18 = functools.partial(ResnetBuilder.build_resnet_18, num_outputs=(2,), output_names=('polarity',))
resnet34 = functools.partial(ResnetBuilder.build_resnet_34, num_outputs=(2,), output_names=('polarity',))
resnet50 = functools.partial(ResnetBuilder.build_resnet_50, num_outputs=(2,), output_names=('polarity',))
resnet101 = functools.partial(ResnetBuilder.build_resnet_101, num_outputs=(2,), output_names=('polarity',))
resnet152 = functools.partial(ResnetBuilder.build_resnet_152, num_outputs=(2,), output_names=('polarity',))

# densenet40 = functools.partial(densenet_l2, depth=40)
# densenet22 = functools.partial(densenet_l2, depth=22)
# densenet16 = functools.partial(densenet_l2, depth=16)
